import json
import os
import uuid

from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
    MessagesPlaceholder,
)
from langchain_deepseek import ChatDeepSeek

from utils.EmbeddingUtil import get_text_embedding
from utils.MilvusClient import MilvusClient
from utils.MySQLClient import MySQLClient

load_dotenv()



class GovInfoAccess:

    def __init__(self, prompt_file: str = None):
        self.history_limit = 3
        self.history = []  # 短期对话记忆
        self.mysql = MySQLClient()  # 长期记忆存储
        self.milvus = MilvusClient()  # 向量数据库
        self.main_prompt_file = prompt_file
        self.llm = ChatDeepSeek(
            api_base=os.getenv("DEEPSEEK_API_BASE"),
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            model="deepseek-v3-250324",
            temperature=0,
            max_tokens=None,
            timeout=None,
            max_retries=2
        )
        self.__init_prompt_templates()
        self.__init_chains()
        self.__init_intent_recognizer()

    def __init_intent_recognizer(self):
        """初始化意图识别组件"""
        self.intent_prompt = ChatPromptTemplate.from_messages([
            SystemMessagePromptTemplate.from_template("""
                你是一个政务事项意图识别专家。请分析用户输入，识别以下内容：
                1. 用户要办理的行政服务事项名称（如食品生产许可证）
                2. 业务类型（新办、变更、注销、延续、换照）
                
                输出要求：
                - 用JSON格式返回识别结果
                - 若无法识别请返回null值
                - 示例：{{"service_item": "食品生产许可证", "business_type": "新办"}}
                """),
            HumanMessagePromptTemplate.from_template("{input}")
        ])
        self.intent_chain = self.intent_prompt | self.llm

    def __init_prompt_templates(self):
        with open(self.main_prompt_file, 'r', encoding='utf-8') as f:
            self.chat_prompt_template = ChatPromptTemplate.from_messages([
                SystemMessagePromptTemplate.from_template(f.read()),
                MessagesPlaceholder("history"),
                HumanMessagePromptTemplate.from_template("{input}")
            ])

    def __init_chains(self):
        self.main_chain = (self.chat_prompt_template | self.llm)

    def _recognize_intent(self, input: str) -> dict:
        """执行意图识别（增加空值保护）"""
        try:
            response = self.intent_chain.invoke({"input": input})
            # 添加响应内容验证
            if not response or not response.content:
                return {"service_item": None, "business_type": None}

            # 处理可能的格式错误
            cleaned_content = response.content.strip("`").replace("json", "", 1)
            return json.loads(cleaned_content) if cleaned_content else {"service_item": None, "business_type": None}
        except Exception as e:
            print(f"意图识别失败: {str(e)}")
            return {"service_item": None, "business_type": None}  # 确保始终返回有效字典

    def _store_conversation(self, query, response, session_id=None):
        """存储对话到MySQL数据库"""
        session_id = session_id or str(uuid.uuid4())
        self.mysql.insert_data(
            table="chat_history",
            data={
                "query": query,
                "response": response,
                "session_id": session_id
            }
        )
        return session_id

    def chat(self, input: str, session_id: str = None) -> tuple:
        """
        多轮对话处理
        :param session_id: 会话ID
        :param input: 用户输入
        :return: AI响应
        """
        # 1. 检索相关历史对话
        # 获取长期上下文（完整会话历史）
        if self.history:
            long_term_context = self._get_long_term_context(session_id)
            # 组合长短期上下文
            combined_context = long_term_context + self.history
        else:
            combined_context = self.history

        # 2. 构建对话历史
        context = []
        for result in combined_context:  # 假设返回相似历史
            # 处理元组格式的历史记录
            if isinstance(result, tuple):
                query_msg, response_msg = result
                context.extend([
                    HumanMessage(content=query_msg),
                    AIMessage(content=response_msg)
                ])
            # 处理字典格式的历史记录
            else:
                context.extend([
                    HumanMessage(content=result.get("query")),
                    AIMessage(content=result.get("response"))
                ])

        # 3. 生成响应
        response = ""
        context.extend([HumanMessage(content=input)])
        for chunk in self.main_chain.stream(
                {'history': context, 'input': input, 'content': "我在中国南昌，今天是2025年4月17日 星期四"}):
            response += chunk.text()

        new_session_id = self._update_storages(input, response, session_id)

        return response, new_session_id

    def stream_chat(self, input: str, session_id: str = None) -> tuple:
        """
        多轮对话处理
        :param session_id: 会话ID
        :param input: 用户输入
        :return: AI响应
        """
        # 检索相关历史对话
        # 获取长期上下文（完整会话历史）
        if self.history:
            long_term_context = self._get_long_term_context(session_id)
            # 组合长短期上下文
            combined_context = long_term_context + self.history
        else:
            combined_context = self.history

        # 构建对话历史
        history = []
        for result in combined_context:  # 假设返回相似历史
            # 处理元组格式的历史记录
            if isinstance(result, tuple):
                query_msg, response_msg = result
                history.extend([
                    HumanMessage(content=query_msg),
                    AIMessage(content=response_msg)
                ])
            # 处理字典格式的历史记录
            else:
                history.extend([
                    HumanMessage(content=result.get("query")),
                    AIMessage(content=result.get("response"))
                ])

        # 新增意图识别步骤
        # 修改意图识别后的处理逻辑
        intent = self._recognize_intent(input) or {}  # 添加空字典保护
        print(f"识别到意图：{intent}")

        # 修改后的检索逻辑
        if intent.get("service_item"):
            try:
                # 获取向量时添加维度转换
                query_vector = get_text_embedding(f"{intent['service_item']}-{intent['business_type']}")
                # 确保是二维列表格式
                milvus_results = self.milvus.search(
                    collection=self.milvus.get_collection("government_items"),
                    query_vectors=[query_vector],  # 注意这里包裹成二维列表
                    top_k=5
                )
                # 处理结果
                if milvus_results:
                    item_code = milvus_results[0].get("item_code")
                    item = self.mysql.get_item_content(item_code)
                    context = item.get("content", "未找到相关内容")
            except Exception as e:
                print(f"Milvus检索异常: {e}")
                context = "Milvus检索过程中发生错误"

        else:
            context = "未识别到用户意图，无法提供相关服务。请重新描述您的问题。"

        # 生成响应
        response_buffer = ""
        history.extend([HumanMessage(content=input)])
        for chunk in self.main_chain.stream(
                {'history': history, 'input': input, 'content': context}):
            response_buffer += chunk.text()
            yield chunk.text()

        self._update_storages(input, response_buffer, session_id)


    def _get_long_term_context(self, session_id: str) -> list:
        """从MySQL获取完整会话历史"""
        if not session_id:
            return []

        return self.mysql.search_data(
            table="chat_history",
            conditions={"session_id": session_id},
            limit=1000  # 获取完整会话历史
        )

    def _update_storages(self, query: str, response: str, session_id: str):
        """更新存储系统"""
        # 长期存储：持久化到MySQL（全量存储）
        new_session_id = session_id or str(uuid.uuid4())
        self.mysql.insert_data(
            table="chat_history",
            data={
                "session_id": new_session_id,
                "query": query,
                "response": response
            }
        )

        # 短期记忆：维护最近对话（内存缓存）
        self.history.append((query, response))
        if len(self.history) > self.history_limit:
            self.history.pop(0)

        return new_session_id


# 使用示例
if __name__ == "__main__":
    access = GovInfoAccess()
    current_dialog_id = None  # 初始对话ID

    # 模拟多轮对话
    dialogues = [
        "食品生产许可证 怎么办的啊"
    ]

    for query in dialogues:
        print(f"用户: {query}")
        response, current_dialog_id = access.chat(query, session_id =current_dialog_id)
        print(f"AI: {response}\n当前会话ID: {current_dialog_id}\n")
